from http import HTTPStatus
import os
from vlmeval.api.base import BaseAPI
from vlmeval.smp import *


# Note: This is a pure language model API.
class QwenAPI(BaseAPI):

    is_api: bool = True

    def __init__(
        self,
        model: str = "qwen-max-1201",
        retry: int = 5,
        wait: int = 5,
        verbose: bool = True,
        seed: int = 2680,
        temperature: float = 0.0,
        system_prompt: str = None,
        key: str = None,
        max_tokens: int = 1024,
        proxy: str = None,
        **kwargs,
    ):

        assert model in [
            "qwen-turbo",
            "qwen-plus",
            "qwen-max",
            "qwen-max-1201",
            "qwen-max-longcontext",
        ]
        self.model = model
        import dashscope

        self.fail_msg = "Failed to obtain answer via API. "
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.seed = seed
        if key is None:
            key = os.environ.get("DASHSCOPE_API_KEY", None)
        assert key is not None, (
            "Please set the API Key (obtain it here: "
            "https://help.aliyun.com/zh/dashscope/developer-reference/vl-plus-quick-start)"
        )
        dashscope.api_key = key
        if proxy is not None:
            proxy_set(proxy)
        super().__init__(
            wait=wait,
            retry=retry,
            system_prompt=system_prompt,
            verbose=verbose,
            **kwargs,
        )

    @staticmethod
    def build_msgs(msgs_raw, system_prompt=None):
        msgs = cp.deepcopy(msgs_raw)
        ret = []
        if system_prompt is not None:
            ret.append(dict(role="system", content=system_prompt))
        for i, msg in enumerate(msgs):
            role = "user" if i % 2 == 0 else "assistant"
            ret.append(dict(role=role, content=msg))
        return ret

    def generate_inner(self, inputs, **kwargs) -> str:
        from dashscope import MultiModalConversation

        assert isinstance(inputs, str) or isinstance(inputs, list)
        inputs = [inputs] if isinstance(inputs, str) else inputs
        messages = self.build_msgs(msgs_raw=inputs, system_prompt=self.system_prompt)

        import dashscope

        response = dashscope.Generation.call(
            model=self.model,
            messages=messages,
            seed=self.seed,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            result_format="message",  # set the result to be "message" format.
        )
        if response.status_code != HTTPStatus.OK:
            return (
                -1,
                "Error: Bad Response Statuse Code. ",
                f"The response status code is {response.status_code}. ",
            )

        try:
            return (
                0,
                response["output"]["choices"][0]["message"]["content"].strip(),
                "Succeeded! ",
            )
        except Exception as err:
            return -1, f"Error: Failed to parse the response. {err}", response
